package com.craftworks.rtmp;

import cn.hutool.core.thread.ThreadUtil;
import com.craftworks.rtmp.amf0.ObjectProperty;
import com.craftworks.rtmp.amf0.ObjectType;
import com.craftworks.rtmp.amf0.StringType;
import org.pcap4j.core.BpfProgram;
import org.pcap4j.core.PcapHandle;
import org.pcap4j.core.PcapNetworkInterface;
import org.pcap4j.core.Pcaps;
import org.pcap4j.packet.Packet;
import org.pcap4j.packet.TcpPacket;

import java.net.NetworkInterface;
import java.util.*;

public class RtmpFilter
{
    private static final RtmpFilter INSTANCE = new RtmpFilter();

    private final Map<Integer, RtmpMessage> cache = new HashMap<>();
    private volatile boolean stopFiltering = true;
    private RtmpFilterListener listener = null;
    private NetworkInterface networkInterface;

    public static RtmpFilter getInstance()
    {
        return INSTANCE;
    }

    public void setListener(RtmpFilterListener listener)
    {
        this.listener = Objects.requireNonNull(listener);
    }

    public void setNetworkInterface(NetworkInterface networkInterface)
    {
        this.networkInterface = Objects.requireNonNull(networkInterface);
    }

    public void start()
    {
        if (networkInterface == null)
            throw new IllegalStateException("network interface not set yet");

        stopFiltering = false;
        ThreadUtil.execute(this::doFiltering);
    }

    public void stop()
    {
        stopFiltering = true;
    }

    private void doFiltering()
    {
        try
        {
            PcapNetworkInterface nif = Pcaps.getDevByAddress(networkInterface.getInetAddresses().nextElement());
            int snapLen = 65536;
            PcapNetworkInterface.PromiscuousMode mode = PcapNetworkInterface.PromiscuousMode.PROMISCUOUS;
            int timeout = 10000;
            cache.clear();
            try (PcapHandle handle = nif.openLive(snapLen, mode, timeout))
            {
                if (listener != null)
                    listener.filterStarted();

                handle.setFilter("tcp", BpfProgram.BpfCompileMode.OPTIMIZE);

                Rtmp.setChunkSize(128);
                RtmpContext ctx = new RtmpContext();
                ctx.csid = -1;
                while (!stopFiltering)
                {
                    Packet packet = handle.getNextPacket();
                    if (packet == null)
                        continue;
                    TcpPacket tcpPacket = packet.getPayload().get(TcpPacket.class);
                    if (tcpPacket.getHeader().getDstPort().valueAsInt() == 1935 &&
                        tcpPacket.getPayload() != null)
                    {
                        boolean done = parseRtmpUrl(tcpPacket.getPayload().getRawData(), ctx);
                        if (done)
                        {
                            if (listener != null)
                            {
                                String rtmpUrl = ctx.connectDomain + "/" + ctx.playOption;
                                listener.urlFiltered(rtmpUrl);
                            }
                            ctx.connectDomain = null;
                            ctx.playOption = null;
                            Rtmp.setChunkSize(128);
                        }
                    }
                }
            }
        } catch (Exception ex)
        {
            stopFiltering = true;
        }

        if (listener != null)
            listener.filterStopped();
    }

    private boolean parseRtmpUrl(byte[] payload, RtmpContext context)
    {
        try
        {
            List<Chunk> chunks = Rtmp.decode(payload);

            List<RtmpMessage> messages = new ArrayList<>();
            RtmpMessage msg;
            for (Chunk chunk : chunks)
            {
                if (chunk.getFormat() == 0 && chunk.getMessageTypeID() == 0x14)
                {
                    msg = RtmpMessage.fromChunk(chunk);
                    if (msg.isComplete())
                    {
                        messages.add(msg);
                        context.csid = -1;
                    } else
                    {
                        cache.put(chunk.getChunkStreamID(), msg);
                        context.csid = chunk.getChunkStreamID();
                    }
                    continue;
                }

                if (chunk.getFormat() != 0)
                {
                    msg = cache.get(context.csid);
                    if (msg != null && chunk.getChunkStreamID() == context.csid)
                    {
                        msg.fillChunk(chunk);
                        if (msg.isComplete())
                        {
                            messages.add(msg);
                            context.csid = -1;
                        }
                    }
                }
            }

            for (RtmpMessage message : messages)
            {
                RtmpBody body = new RtmpBody(message.getBody());
                StringType cmdName = (StringType) body.getDataTypeList().get(0);
                if (cmdName.stringValue().equals("connect"))
                {
                    ObjectType objectType = (ObjectType) body.getDataTypeList().get(2);
                    Map<String, ObjectProperty> properties = objectType.propertyValue();
                    ObjectProperty tcUrl = properties.get("tcUrl");
                    context.connectDomain = tcUrl.getValue().stringValue();
                    context.playOption = null;
                    context.csid = -1;
                    return false;
                }
                if (context.connectDomain != null && cmdName.stringValue().equals("play"))
                {
                    StringType option = (StringType) body.getDataTypeList().get(3);
                    context.playOption = option.stringValue();
                    context.csid = -1;
                    return true;
                }
            }
            return false;
        } catch (Exception ex)
        {
            return false;
        }
    }

    public boolean isStarted()
    {
        return !stopFiltering;
    }

    static class RtmpContext
    {
        int csid;
        String connectDomain;
        String playOption;
    }
}
